# 导入所需的库
import code
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback
import os
import logging
import gc
from datasets import load_dataset
import torch
import multiprocessing

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,  # 设置日志级别为INFO
    format='%(asctime)s - %(levelname)s - %(message)s',  # 设置日志格式
    handlers=[
        logging.FileHandler("training.log", encoding='utf-8'),  # 将日志写入文件，指定UTF-8编码
        logging.StreamHandler()  # 将日志输出到控制台
    ]
)
logger = logging.getLogger(__name__)  # 获取日志记录器

# 定义预训练模型的路径
model_path = "./deepseek-ai/Deepseek-R1-Distill-Qwen-1___5B"
logger.info(f"开始加载tokenizer: {model_path}")
try:
    # 从指定路径加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    # 设置填充标记为结束标记，这对于某些模型是必要的
    tokenizer.pad_token = tokenizer.eos_token
    logger.info("Tokenizer加载成功")
except Exception as e:
    logger.error(f"加载tokenizer时出错: {e}")
    raise  # 如果加载失败，抛出异常

# 检测可用的计算设备（优先使用CUDA GPU，否则使用CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"使用设备: {device}")

# 手动触发垃圾回收以释放内存
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()  # 如果使用CUDA，清空CUDA缓存

# 加载预训练模型
logger.info("开始加载模型...")
try:
    model = AutoModelForCausalLM.from_pretrained(
        model_path,  # 模型路径
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # 根据设备选择数据类型（GPU使用float16以节省内存，CPU使用float32）
        device_map="auto",  # 自动将模型分片加载到可用设备上
        low_cpu_mem_usage=True  # 优化CPU内存使用
    )
    logger.info("模型加载成功")
except Exception as e:
    logger.error(f"加载模型时出错: {e}")
    raise  # 如果加载失败，抛出异常

# 加载数据集
try:
    logger.info("开始加载数据集...")
    # 检查数据文件是否存在
    if not os.path.exists("./tiangou.json"):
        logger.error("数据文件 ./tiangou.json 不存在")
        raise FileNotFoundError("数据文件 ./tiangou.json 不存在")
        
    # 从JSON文件加载数据集
    dataset = load_dataset("json", data_files="./tiangou.json")
    logger.info(f"数据集加载完成，样本数量: {len(dataset['train'])}")
    
    # 检查数据集是否为空
    if len(dataset['train']) == 0:
        logger.warning("数据集为空，请检查数据文件")
    
    # 检查数据集中是否包含必需的'conversation'字段
    if "conversation" not in dataset['train'].column_names:
        logger.error("数据集中缺少conversation字段，请检查数据格式")
        raise ValueError("数据集中缺少conversation字段")
        
    # 打印数据集的第一个样本以供调试
    logger.info(f"数据集样本示例: {dataset['train'][0] if len(dataset['train']) > 0 else 'No samples'}")
    
except Exception as e:
    logger.error(f"加载数据集时出错: {e}")
    raise  # 如果加载失败，抛出异常

# 定义数据预处理函数
def preprocess_function(examples):
    """将原始数据集转换为模型可以理解的格式。"""
    try:
        inputs = []  # 存储处理后的输入文本
        targets = []  # 存储处理后的目标文本
        
        # 遍历数据集中的每个样本
        for i, conversation_list in enumerate(examples["conversation"]):
            # 确保'conversation'字段是列表类型
            if not isinstance(conversation_list, list):
                logger.warning(f"样本 {i} 的conversation不是列表类型，跳过此样本: {conversation_list}")
                continue

            # 提取系统提示（如果存在）
            system_prompt = ""
            for turn in conversation_list:
                if isinstance(turn, dict) and turn.get("system"):
                    system_prompt = turn.get("system")
                    break
            
            # 遍历对话列表中的每一项（可能是包含多轮对话的列表）
            for turn in conversation_list:
                # 确保每一轮是字典
                if not isinstance(turn, dict):
                    logger.warning(f"样本 {i} 中的对话轮次不是字典类型，跳过此轮: {turn}")
                    continue

                # 提取有效的 input 和 output
                input_text = turn.get("input")
                output_text = turn.get("output")

                # 确保 input 和 output 都存在且不为空
                if input_text and output_text:
                    # 构建完整的提示词格式
                    if system_prompt:
                        # 如果有系统提示，将其添加到输入中
                        formatted_input = f"<system>\n{system_prompt}\n</system>\n\n<user>\n{input_text}\n</user>\n\n<assistant>\n"
                    else:
                        # 如果没有系统提示，只使用用户输入
                        formatted_input = f"<user>\n{input_text}\n</user>\n\n<assistant>\n"
                    
                    # 格式化输出（确保模型学习生成完整的回复）
                    formatted_output = f"{output_text}\n</assistant>"
                    
                    inputs.append(formatted_input)
                    targets.append(formatted_output)
                    
                    logger.info(f"样本 {i} 处理后的输入: {formatted_input[:50]}...")
                    logger.info(f"样本 {i} 处理后的输出: {formatted_output[:50]}...")
                elif output_text and not input_text:
                    # 处理只有 output 的情况
                    logger.warning(f"样本 {i} 轮次缺少 input: {turn}")
                elif input_text and not output_text:
                    # 处理只有 input 的情况
                    logger.warning(f"样本 {i} 轮次缺少 output: {turn}")
        
        # 使用分词器对输入和目标文本进行编码
        # 增加最大长度以容纳更长的提示词格式
        model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
        labels = tokenizer(targets, max_length=512, truncation=True, padding="max_length")
        
        # 将目标文本的input_ids作为模型的标签
        model_inputs["labels"] = labels["input_ids"].copy()
        
        # 将标签中对应填充标记的部分设置为-100，这样在计算损失时会被忽略
        for i in range(len(model_inputs["labels"])):
            for j in range(len(model_inputs["labels"][i])):
                if labels["input_ids"][i][j] == tokenizer.pad_token_id:
                    model_inputs["labels"][i][j] = -100
        
        return model_inputs  # 返回处理后的模型输入
    except Exception as e:
        logger.error(f"预处理数据时出错: {e}")
        # 返回空字典以避免整个过程失败
        return {"input_ids": [], "attention_mask": [], "labels": []}

# 定义应用预处理函数的函数
def preprocess_data(dataset, tokenizer):
    """对整个数据集应用预处理函数。"""
    logger.info("开始预处理数据...")

    # 清理内存
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    try:
        # 使用.map()方法对数据集进行批处理
        # batched=True: 以批次方式处理数据
        # remove_columns: 移除原始的'conversation'列
        # num_proc=1: 使用单进程处理，避免多进程可能导致的问题
        # desc: 进度条描述
        # load_from_cache_file=False: 禁用缓存，确保每次都重新处理
        tokenized_dataset = dataset.map(
            preprocess_function, 
            batched=True, 
            remove_columns=["conversation"],
            num_proc=1,  
            desc="预处理训练数据",
            load_from_cache_file=True # 启用缓存以加速后续运行  
        )
        logger.info(f"数据预处理完成，处理后样本数量: {len(tokenized_dataset['train'])}")
        
        # 检查处理后的数据集是否为空
        if len(tokenized_dataset['train']) == 0:
            logger.error("预处理后的数据集为空，请检查预处理函数")
            raise ValueError("预处理后的数据集为空")
        
        # 打印一些统计信息以验证预处理
        input_lengths = [len(x) for x in tokenized_dataset["train"]["input_ids"]]
        logger.info(f"输入序列长度统计: 最小={min(input_lengths)}, 最大={max(input_lengths)}, 平均={sum(input_lengths)/len(input_lengths):.2f}")
        
        # 检查非填充标记的数量
        non_pad_tokens = sum([sum([1 for token in sample if token != tokenizer.pad_token_id]) for sample in tokenized_dataset["train"]["input_ids"]])
        logger.info(f"非填充标记总数: {non_pad_tokens}")
        
        # 打印一个处理后的样本示例
        if len(tokenized_dataset['train']) > 0:
            logger.info(f"处理后的样本示例: {tokenized_dataset['train'][0]}")
            
        return tokenized_dataset  # 返回处理后的数据集
    except Exception as e:
        logger.error(f"预处理数据时出错: {e}")
        raise  # 如果预处理失败，抛出异常

# 定义设置训练参数的函数
def setup_training_args():
    """配置Hugging Face Trainer所需的训练参数。"""
    logger.info("设置训练参数...")
    try:
        # 创建输出和日志目录
        os.makedirs("./results", exist_ok=True)
        os.makedirs("./logs", exist_ok=True)
        
        training_args = TrainingArguments(
            output_dir="./results",  # 模型检查点和输出的保存目录
            num_train_epochs=10,  # 增加训练轮数以确保充分学习
            per_device_train_batch_size=1,  # 减小批次大小以提高稳定性
            gradient_accumulation_steps=16,  # 增加梯度累积步数以保持有效批次大小
            save_strategy="steps",  # 按步数保存检查点
            save_steps=200,  # 更频繁地保存检查点
            save_total_limit=3,  # 最多保存3个检查点
            logging_dir="./logs",  # 日志保存目录
            logging_strategy="steps",  # 按步数记录日志
            logging_steps=20,  # 更频繁地记录日志
            learning_rate=5e-6,  # 降低学习率以提高稳定性
            warmup_ratio=0.2,  # 增加预热比例
            fp16=torch.cuda.is_available(),  # 如果有CUDA GPU，则启用混合精度训练
            fp16_full_eval=False,  # 禁用全精度评估以节省内存
            report_to="none",  # 不向外部服务报告训练指标
            dataloader_num_workers=0,  # 数据加载器的工作进程数（设为0避免多进程问题）
            dataloader_pin_memory=True,  # 锁定内存以加速数据传输
            weight_decay=0.01,  # 权重衰减（L2正则化）
            max_grad_norm=0.5,  # 梯度裁剪阈值，防止梯度爆炸
            remove_unused_columns=True,  # 自动移除数据集中模型不需要的列
            disable_tqdm=False,  # 显示训练进度条
            # 移除评估相关参数，因为没有提供验证集
            optim="adamw_torch",  # 使用PyTorch实现的AdamW优化器
            no_cuda=True,  # 告知Trainer不要尝试自动将模型移至CUDA（因为device_map="auto"已处理）
          )
        logger.info("训练参数设置完成")
        return training_args
    except Exception as e:
        logger.error(f"设置训练参数时出错: {e}")
        raise  # 如果设置失败，抛出异常

# 创建自定义回调以监控内存使用情况
class MemoryMonitorCallback(TrainerCallback):
    """一个自定义的回调函数，用于在训练过程中监控GPU内存使用情况。"""
    def on_step_end(self, args, state, control, **kwargs):
        # 每50步检查一次内存
        if state.global_step % 50 == 0:  
            if torch.cuda.is_available():
                memory_allocated = torch.cuda.memory_allocated() / 1024**2  # 已分配内存（MB）
                memory_reserved = torch.cuda.memory_reserved() / 1024**2  # 已预留内存（MB）
                logger.info(f"Step {state.global_step}: CUDA Memory Allocated: {memory_allocated:.2f} MB, Reserved: {memory_reserved:.2f} MB")
                # 如果内存使用接近上限，尝试清理
                if memory_allocated > 0.9 * torch.cuda.get_device_properties(0).total_memory / 1024**2:
                    logger.warning("内存使用过高，尝试清理...")
                    gc.collect()  # 垃圾回收
                    torch.cuda.empty_cache()  # 清空CUDA缓存

# 定义创建Trainer和训练模型的函数
def train_model(model, tokenizer, tokenized_dataset, training_args):
    """初始化Trainer并开始模型训练。"""
    logger.info("创建Trainer...")
    try:
        # 清理内存
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
        # 初始化Trainer
        trainer = Trainer(
            model=model,  # 要训练的模型
            args=training_args,  # 训练参数
            train_dataset=tokenized_dataset["train"],  # 训练数据集
            callbacks=[MemoryMonitorCallback()],  # 添加内存监控回调
            # 不需要显式设置数据整理器（data_collator），Trainer会使用默认的
        )
        logger.info("Trainer创建成功")
        
        # 开始训练
        logger.info("开始训练模型...")
        # 清理内存
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
        trainer.train()  # 执行训练过程
        logger.info("训练完成")
        
        # 保存微调后的模型
        logger.info("训练完成，正在保存模型...")
        # 创建保存目录
        os.makedirs("./fine_tuned_model", exist_ok=True)
        
        # 保存模型权重和配置
        model.save_pretrained("./fine_tuned_model")
        # 保存分词器
        tokenizer.save_pretrained("./fine_tuned_model")
        logger.info("模型保存完成")
        
    except Exception as e:
        logger.error(f"训练过程中出错: {e}")
        raise  # 如果训练失败，抛出异常

# 主函数，协调整个流程
def main():
    """执行模型微调的主要流程。"""
    # 加载模型和tokenizer (这部分在脚本开头已经完成，但在main函数中重复以便独立运行)
    model_path = "./deepseek-ai/Deepseek-R1-Distill-Qwen-1___5B"
    logger.info(f"开始加载tokenizer: {model_path}")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        tokenizer.pad_token = tokenizer.eos_token
        logger.info("Tokenizer加载成功")
    except Exception as e:
        logger.error(f"加载tokenizer时出错: {e}")
        raise

    # 检测设备 (重复)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"使用设备: {device}")

    # 清理内存 (重复)
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # 加载模型 (重复)
    logger.info("开始加载模型...")
    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto",
            low_cpu_mem_usage=True
        )
        logger.info("模型加载成功")
    except Exception as e:
        logger.error(f"加载模型时出错: {e}")
        raise

    # 加载数据集 (重复)
    try:
        logger.info("开始加载数据集...")
        if not os.path.exists("./tiangou.json"):
            logger.error("数据文件 ./tiangou.json 不存在")
            raise FileNotFoundError("数据文件 ./tiangou.json 不存在")
            
        dataset = load_dataset("json", data_files="./tiangou.json")
        logger.info(f"数据集加载完成，样本数量: {len(dataset['train'])}")
        
        if len(dataset['train']) == 0:
            logger.warning("数据集为空，请检查数据文件")
        
        if "conversation" not in dataset['train'].column_names:
            logger.error("数据集中缺少conversation字段，请检查数据格式")
            raise ValueError("数据集中缺少conversation字段")
            
        logger.info(f"数据集样本示例: {dataset['train'][0] if len(dataset['train']) > 0 else 'No samples'}")
        
    except Exception as e:
        logger.error(f"加载数据集时出错: {e}")
        raise
    
    # 预处理数据
    tokenized_dataset = preprocess_data(dataset, tokenizer)
    
    # 设置训练参数
    training_args = setup_training_args()
    
    # 训练模型
    train_model(model, tokenizer, tokenized_dataset, training_args)

# Python脚本的主入口点
if __name__ == '__main__':
    # 添加这行是为了在Windows等多进程环境中正确启动子进程
    multiprocessing.freeze_support()
    main()  # 调用主函数开始执行